# import 2D backbone
from .backbone_2d.yolov2 import build_yolov2
# import 3D backbone
from .backbone_3d.shufflnetv2 import build_shufflenetv2_3d

__all__ = ['build_backbone_2d', 'build_backbone_3d']
def build_backbone_2d(cfg, pretrained=False):
    print('==============================')
    print('2D Backbone: {}'.format(cfg['backbone_2d'].upper()))
    print('--pretrained: {}'.format(pretrained))

    model, feat_dim = build_yolov2(pretrained)

    return model, feat_dim


def build_backbone_3d(cfg, pretrained=False):
    print('==============================')
    print('3D Backbone: {}'.format(cfg['backbone_3d'].upper()))
    print('--pretrained: {}'.format(pretrained))

    model, feat_dim = build_shufflenetv2_3d(
        model_size=cfg['model_size'],
        pretrained=pretrained
        )

    return model, feat_dim
